from json.tool import main
import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler, LabelEncoder
from sklearn.metrics import accuracy_score, classification_report
from sklearn.neighbors import KNeighborsClassifier
from sklearn.svm import SVC
from sklearn.linear_model import LogisticRegression
from sklearn.tree import DecisionTreeClassifier
from sklearn.ensemble import RandomForestClassifier
from joblib import dump

class Classification:  # 修正类名拼写错误（原Classfication）
    def __init__(self):
        self.df = pd.read_csv('fina_indicator.csv')

    def get_conditions(self):
        # 分类前准备
        df = self.df
        # 计算收益和风险比例
        df['max_ratio'] = df['max_close'] / df['the_close']  # 收益潜力
        df['min_ratio'] = df['min_close'] / df['the_close']  # 风险潜力

        # 自动划分高低阈值
        high_return_threshold = df['max_ratio'].quantile(0.4)  # 前40%定义为高收益
        high_risk_threshold = df['min_ratio'].quantile(0.4)  # 前40%定义为高风险

        # 生成分类标签
        conditions = [
            (df['max_ratio'] >= high_return_threshold) & (df['min_ratio'] >= high_risk_threshold),
            (df['max_ratio'] < high_return_threshold) & (df['min_ratio'] >= high_risk_threshold),
            (df['max_ratio'] >= high_return_threshold) & (df['min_ratio'] < high_risk_threshold),
            (df['max_ratio'] < high_return_threshold) & (df['min_ratio'] < high_risk_threshold),
        ]
        choices = ['高收益高风险', '低收益高风险', '高收益低风险', '低收益低风险']
        df['category'] = np.select(conditions, choices, default='未知')
        
        features = ['eps', 'total_revenue_ps', 'undist_profit_ps', 'gross_margin', 'fcff', 'fcfe', 'tangible_asset', 'bps',
                               'grossprofit_margin', 'npta']
        # 提取特征矩阵（修正原代码中直接使用features列表的问题）
        X_features = df[features].values
        
        label_encoder = LabelEncoder()
        df['category_encoded'] = label_encoder.fit_transform(df['category'])
        scaler = StandardScaler()
        X = scaler.fit_transform(X_features)  # 使用特征矩阵进行标准化
        y = df['category_encoded']
        
        # 划分训练集和测试集
        X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=42)
        dump(scaler, 'scaler_joblib')
        return X_train, X_test, y_train, y_test, label_encoder

    def knn_utils(self, X_train, X_test, y_train, y_test, label_encoder):
        """KNN分类器"""
        # 初始化KNN分类器
        knn = KNeighborsClassifier(n_neighbors=5)
        # 训练模型
        knn.fit(X_train, y_train)
        # 预测
        y_pred = knn.predict(X_test)
        # 评估模型
        accuracy = accuracy_score(y_test, y_pred)
        report = classification_report(y_test, y_pred, target_names=label_encoder.classes_)
        print(f"KNN Accuracy: {accuracy:.4f}")
        print("KNN Classification Report:")
        print(report)
        dump(knn, 'knn_model_joblib')
        dump(label_encoder, 'label_encoder_joblib')

    def svc_utils(self, X_train, X_test, y_train, y_test, label_encoder):  # 修正缩进和方法定义
        """SVC分类器"""
        # 初始化SVC分类器
        svc = SVC(kernel='linear', C=1.0, random_state=42)
        # 训练模型
        svc.fit(X_train, y_train)
        # 预测
        y_pred = svc.predict(X_test)
        # 评估模型
        print("SVC Classification Report:")
        print(classification_report(y_test, y_pred, target_names=label_encoder.classes_))
        dump(svc, 'svc_model.joblib')
        dump(label_encoder, 'label_encoder.joblib')

    def logistic_regression_utils(self, X_train, X_test, y_train, y_test, label_encoder):
        """逻辑回归分类器"""
        # 初始化逻辑回归模型
        lr = LogisticRegression(max_iter=1000, random_state=42)  # 增加max_iter避免收敛警告
        # 训练模型
        lr.fit(X_train, y_train)
        # 预测
        y_pred = lr.predict(X_test)
        # 评估模型
        accuracy = accuracy_score(y_test, y_pred)
        print(f"Logistic Regression Accuracy: {accuracy:.4f}")
        print("Logistic Regression Classification Report:")
        print(classification_report(y_test, y_pred, target_names=label_encoder.classes_))
        dump(lr, 'logistic_regression_model.joblib')
        dump(label_encoder, 'label_encoder_joblib')

    def decision_tree_utils(self, X_train, X_test, y_train, y_test, label_encoder):
        """决策树分类器"""
        # 初始化决策树模型
        dt = DecisionTreeClassifier(max_depth=5, random_state=42)  # 限制树深防止过拟合
        # 训练模型
        dt.fit(X_train, y_train)
        # 预测
        y_pred = dt.predict(X_test)
        # 评估模型
        accuracy = accuracy_score(y_test, y_pred)
        print(f"Decision Tree Accuracy: {accuracy:.4f}")
        print("Decision Tree Classification Report:")
        print(classification_report(y_test, y_pred, target_names=label_encoder.classes_))
        dump(dt, 'decision_tree_model.joblib')
        dump(label_encoder, 'label_encoder_joblib')

    def random_forest_utils(self, X_train, X_test, y_train, y_test, label_encoder):
        """随机森林分类器"""
        # 初始化随机森林模型
        rf = RandomForestClassifier(n_estimators=100, max_depth=5, random_state=42)
        # 训练模型
        rf.fit(X_train, y_train)
        # 预测
        y_pred = rf.predict(X_test)
        # 评估模型
        accuracy = accuracy_score(y_test, y_pred)
        print(f"Random Forest Accuracy: {accuracy:.4f}")
        print("Random Forest Classification Report:")
        print(classification_report(y_test, y_pred, target_names=label_encoder.classes_))
        dump(rf, 'random_forest_model.joblib')
        dump(label_encoder, 'label_encoder_joblib')


if __name__ == '__main__':
    clf = Classification()
    X_train, X_test, y_train, y_test, label_encoder = clf.get_conditions()
    
    # 可选择运行不同的分类器
    clf.knn_utils(X_train, X_test, y_train, y_test, label_encoder)
    clf.svc_utils(X_train, X_test, y_train, y_test, label_encoder)
    clf.logistic_regression_utils(X_train, X_test, y_train, y_test, label_encoder)
    clf.decision_tree_utils(X_train, X_test, y_train, y_test, label_encoder)
    clf.random_forest_utils(X_train, X_test, y_train, y_test, label_encoder)
